import asyncio
import datetime
import fractions
import json
import time

import cv2
import numpy as np
import pyautogui
import websockets
from aiortc import RTCPeerConnection, MediaStreamTrack, RTCConfiguration, RTCIceServer, RTCSessionDescription
from aiortc.contrib.media import MediaRelay, MediaPlayer
from av import VideoFrame

from logger import log
from settings import settings

pcs = {}
key_mapping = {
    'ArrowUp': 'up',
    'ArrowDown': 'down',
    'ArrowLeft': 'left',
    'ArrowRight': 'right'
}


class VideoStreamTrack(MediaStreamTrack):
    kind = "video"

    def __init__(self):
        super().__init__()  # 初始化父类
        self._frame = None
        self._start_time = time.time()

    async def recv(self):
        # 假设我们想要30fps的帧率
        fps = 10
        time_base = 1 / fps

        # 计算当前帧的 pts
        current_time = time.time()
        elapsed_time = current_time - self._start_time
        pts = int(elapsed_time * fps)

        # 在这里捕获桌面截图并转换为 VideoFrame
        img = pyautogui.screenshot()
        frame = np.array(img)[:, :, :3]  # 删除alpha通道
        frame = cv2.resize(frame, (1280, 720))
        # 将 RGB 格式转换为 BGR 格式
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        video_frame = VideoFrame.from_ndarray(frame, format='bgr24')
        video_frame.pts = pts
        video_frame.time_base = fractions.Fraction(1, fps)

        await asyncio.sleep(time_base)  # 控制帧率
        return video_frame


async def onOffer(ws, params):
    log.info("offer")
    pc = RTCPeerConnection(RTCConfiguration([RTCIceServer("turn:39.106.154.90:3478", "lipeng", "123456")]))
    # pc = RTCPeerConnection()
    from_ = params["from"]
    pcs[from_] = {}
    pcs[from_]['pc'] = pc

    @pc.on("connectionstatechange")
    async def on_connectionstatechange():
        log.info(f"Connection state is {pc.connectionState} for {from_}")
        if pc.connectionState == "failed":
            await pcs[from_]['pc'].close()
            pcs[from_]['relay'] = None
            if pcs[from_]['webcam']:
                pcs[from_]['webcam'].stop()
            pcs[from_]['webcam'] = None
            pcs.pop(from_)
            log.info(f"poped: {pcs}")

    # video_track = MediaPlayer('D:/1.mp4').video
    # 创建视频轨道实例
    video_track = VideoStreamTrack()

    # 使用 MediaRelay 订阅这个轨道以支持多个消费者
    relay = MediaRelay()
    video_relay_track = relay.subscribe(video_track)
    pcs[params['from']]['relay'] = relay
    pcs[params['from']]['webcam'] = video_track

    # 将 relay 订阅的轨道添加到 RTCPeerConnection 中
    pc.addTrack(video_relay_track)

    @pc.on("datachannel")
    def on_datachannel(channel):
        log.info(f"{channel}, -, created by remote party")

        @channel.on("message")
        def on_message(message):
            log.info(f"{channel}<{message}")
            onDatachannel(message)

    await pc.setRemoteDescription(RTCSessionDescription(sdp=params['offer']["sdp"], type=params['offer']["type"]))

    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    answer_rsp = {
        "type": "answer",
        "answer": {"sdp": pc.localDescription.sdp, "type": pc.localDescription.type},
        "from": settings.code,
        "to": from_
    }
    log.info(f"Answer Response")
    await ws.send(json.dumps(answer_rsp))


def onDatachannel(message):
    command = json.loads(message)
    try:
        if command['type'] == 'mousemove_event':
            # 根据视频分辨率和屏幕分辨率计算实际点击位置
            width, height = pyautogui.size()
            x1 = command['position']['x'] * width
            y1 = command['position']['y'] * height
            pyautogui.mouseDown(x=x1, y=y1, button=command['button'])
            x = command['positionUp']['x'] * width
            y = command['positionUp']['y'] * height
            pyautogui.moveTo(x=x, y=y)
            pyautogui.mouseUp(x=x, y=y, button=command['button'])
        if command['type'] == 'mouse_event':
            # 根据视频分辨率和屏幕分辨率计算实际点击位置
            width, height = pyautogui.size()
            screen_x = command['position']['x'] * width
            screen_y = command['position']['y'] * height
            pyautogui.click(x=screen_x, y=screen_y, button=command['button'])
        elif command['type'] == 'keyboard_event':
            key_event = command['event']
            key_value = command['key']
            is_combination = command.get('isCombination', False)

            if key_value in key_mapping:
                pyautogui_key = key_mapping[key_value]
                pyautogui.press(pyautogui_key)
            elif is_combination:
                # 如果是组合键，你可能需要拆分按键并分别处理
                keys = key_value.split('+')
                # 按下所有修饰键
                for key in keys[:-1]:
                    pyautogui.keyDown(key.lower())
                # 按下并释放最后一个键
                pyautogui.press(keys[-1].lower())
                # 释放所有修饰键
                for key in reversed(keys[:-1]):
                    pyautogui.keyUp(key.lower())
            else:
                # 如果是普通键或方向键，直接模拟按键
                pyautogui.press(key_value.lower())
    except Exception as e:
        log.error(f"Error executing command: {e}")


async def onCommand(ws, command):
    log.info(f"onCommand: {command}")


def onPing(ws, loads):
    # log.info("ping")
    global last_timestamp
    last_timestamp = datetime.datetime.now()


async def onMessage(ws, message):
    log.info("Received message:{}".format(message))
    loads = json.loads(message)
    if loads['type'] == 'offer':
        await onOffer(ws, loads)
    elif loads['type'] == 'command':
        await onCommand(ws, loads)
    elif loads['type'] == 'ping':
        onPing(ws, loads)


async def heartbeat(websocket):
    while True:
        # 检测时间差
        await asyncio.sleep(10)  # 每10秒检查一次
        if (datetime.datetime.now() - last_timestamp).seconds > 30:
            # 如果超过30秒没有收到消息，则尝试重连
            log.info("No message in the last 30 seconds, reconnecting...")
            await websocket.close()
            break
        # 发送心跳消息
        try:
            message = {
                "type": "ping",
                "from": settings.code,
                "to": settings.code
            }
            await websocket.send(json.dumps(message))
        except Exception as e:
            log.info(f"An error occurred during heartbeat: {e}")
            break


async def socket_logic():
    # 连接信令服务器
    # uri = "ws://localhost:9443/websocketRTC/" + settings.code
    uri = "wss://webrtc.wechain360.com/websocketRTC/" + settings.code
    # uri = "wss://mq.wechain360.com:8443/websocketRTC/" + settings.code

    while True:
        try:
            async with websockets.connect(uri) as websocket:
                global last_timestamp
                last_timestamp = datetime.datetime.now()
                # 启动心跳任务
                heartbeat_task = asyncio.create_task(heartbeat(websocket))
                log.info("WebSocket connection opened")
                while True:
                    message = await websocket.recv()
                    await onMessage(websocket, message)
        except websockets.exceptions.ConnectionClosed as e:
            log.info(f"WebSocket connection closed with exception: {e}")
        except Exception as e:
            log.info(f"Exception: {e}")
        finally:
            # 等待时间检查任务结束
            if 'heartbeat_task' in locals():
                heartbeat_task.cancel()
                # await heartbeat_task
        await asyncio.sleep(5)  # 等待一段时间后再次尝试连接，避免无限快速重连
